package com.zys.background.auth.util;

import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationCode;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;

import java.util.Map;
import java.util.function.Consumer;

/**
 * @author zys
 * @since 2022-11-21
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
@SuppressWarnings("AlibabaClassNamingShouldBeCamel")
public class OAuth2AuthenticationProviderUtils {

  public static OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(Authentication authentication) {
    OAuth2ClientAuthenticationToken clientPrincipal = null;
    if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(authentication.getPrincipal().getClass())) {
      clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
    }
    if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
      return clientPrincipal;
    }
    throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_CLIENT);
  }

  public static <T extends OAuth2Token> OAuth2Authorization invalidate(OAuth2Authorization authorization, T token) {

    Consumer<Map<String, Object>> metadataConsumer = metadata -> metadata.put(OAuth2Authorization.Token.INVALIDATED_METADATA_NAME, true);

    OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.from(authorization)
        .token(token, metadataConsumer);

    if (OAuth2RefreshToken.class.isAssignableFrom(token.getClass())) {
      authorizationBuilder.token(authorization.getAccessToken().getToken(),metadataConsumer);

      OAuth2Authorization.Token<OAuth2AuthorizationCode> authorizationCode = authorization.getToken(OAuth2AuthorizationCode.class);
      if (authorizationCode != null && !authorizationCode.isInvalidated()) {
        authorizationBuilder.token(authorizationCode.getToken(), metadataConsumer);
      }
    }

    return authorizationBuilder.build();
  }
}
